package network

import (
	"log"
	"strconv"
	"sync"
	"time"
)

// NetworkGRPC implements the network layer using gRPC (WaterBear style)
type NetworkGRPC struct {
	config      NetConfig
	MsgInCh     chan Message
	MsgOutCh    chan Message
	addrs       []string
	msgHandler  func([]byte)
	mu          sync.RWMutex
	started     bool
}

// Ensure NetworkGRPC implements NetworkInterface
var _ NetworkInterface = (*NetworkGRPC)(nil)

// NewNetworkGRPC creates a new gRPC-based network
func NewNetworkGRPC(config NetConfig, msgInCh, msgOutCh chan Message) *NetworkGRPC {
	return &NetworkGRPC{
		config:   config,
		MsgInCh:  msgInCh,
		MsgOutCh: msgOutCh,
		addrs:    config.Addrs,
	}
}

// Start starts the network layer
func (n *NetworkGRPC) Start() {
	n.mu.Lock()
	if n.started {
		n.mu.Unlock()
		return
	}
	n.started = true
	n.mu.Unlock()

	// Start receiver
	port := n.config.Addrs[n.config.ID]
	go n.startReceiver(port)

	// Start sender
	StartSender(int64(n.config.ID), n.addrs)

	// Start message dispatcher
	go n.dispatch()

	log.Printf("Network layer (gRPC) started for Node %d", n.config.ID)
}

// startReceiver starts the gRPC receiver
func (n *NetworkGRPC) startReceiver(port string) {
	// Set up message handler
	n.msgHandler = func(msgBytes []byte) {
		n.handleReceivedMessage(msgBytes)
	}

	// Start receiver
	StartReceiver(strconv.Itoa(n.config.ID), port, n.msgHandler)
}

// handleReceivedMessage handles a received message
func (n *NetworkGRPC) handleReceivedMessage(msgBytes []byte) {
	// First deserialize MessageWithSignature to get the message and MAC
	msgWithSig, err := DeserializeMessageWithSignature(msgBytes)
	if err != nil {
		log.Printf("Node %d: Failed to deserialize message with signature: %v", n.config.ID, err)
		return
	}

	// Verify MAC (simplified - should use proper key)
	senderID := int64(0) // Should extract from message
	if !VerifyMAC(senderID, msgWithSig.Msg, msgWithSig.Sig) {
		log.Printf("Node %d: MAC verification failed", n.config.ID)
		return
	}

	// Deserialize ReplicaMessage
	replicaMsg := DeserializeReplicaMessage(msgWithSig.Msg)

	// Convert to generic Message
	ebbMsg := ConvertReplicaToMessage(replicaMsg, int(replicaMsg.Source))

	// Set sender from the message
	ebbMsg.Sender = int(replicaMsg.Source)

	// Send to input channel
	select {
	case n.MsgInCh <- ebbMsg:
	case <-time.After(5 * time.Second):
		log.Printf("Node %d: WARNING: MsgInCh blocked for 5s", n.config.ID)
		n.MsgInCh <- ebbMsg // Force blocking send
	}
}

// dispatch dispatches messages from MsgOutCh
func (n *NetworkGRPC) dispatch() {
	for msg := range n.MsgOutCh {
		msg.Sender = n.config.ID // Set the sender ID before dispatching

		if msg.To == n.config.ID {
			// Loopback message
			log.Printf("Node %d: Loopback message for self (Round %d, Tag %s)", n.config.ID, msg.Round, msg.Tag)
			n.MsgInCh <- msg
		} else if msg.To == -1 {
			// Broadcast message
			n.Broadcast(msg)
		} else {
			// Point-to-point message
			n.Send(msg.To, msg)
		}
	}
}

// Broadcast broadcasts a message to all peers
func (n *NetworkGRPC) Broadcast(msg Message) {
	// Convert to ReplicaMessage
	replicaMsg := ConvertMessageToReplicaMessage(msg)
	
	// Serialize
	msgBytes, err := replicaMsg.Serialize()
	if err != nil {
		log.Printf("Node %d: Failed to serialize message: %v", n.config.ID, err)
		return
	}

	// Send loopback copy to self
	go func() {
		select {
		case n.MsgInCh <- msg:
		case <-time.After(5 * time.Second):
			log.Printf("Node %d: WARNING: MsgInCh blocked for 5s during broadcast loopback", n.config.ID)
			n.MsgInCh <- msg // Force blocking send
		}
	}()

	// Determine protocol type from message tag
	var mtype ProtocolType
	// CBR protocol uses CBR type
	if msg.Tag == "CBR" || msg.Tag == "CBR_PREPARE" || msg.Tag == "CBR_SHARE_REDUCE" || 
	   msg.Tag == "CBR_PROACTIVIZE" || msg.Tag == "CBR_SHARE_DIST" {
		mtype = CBR
	} else {
		// Default to CBR for unknown tags
		mtype = CBR
	}

	// Broadcast to all nodes
	// Get addresses for all peers
	var peerAddrs []string
	for i, addr := range n.addrs {
		if i != n.config.ID {
			peerAddrs = append(peerAddrs, addr)
		}
	}

	// Use the client factory from grpc_client.go
	MACBroadcast(msgBytes, mtype, peerAddrs, CreateGRPCClient)
}

// Send sends a point-to-point message
func (n *NetworkGRPC) Send(to int, msg Message) {
	if to < 0 || to >= len(n.addrs) {
		log.Printf("Node %d: Invalid target node %d", n.config.ID, to)
		return
	}

	if to == n.config.ID {
		// Loopback
		n.MsgInCh <- msg
		return
	}

	// Convert to ReplicaMessage
	replicaMsg := ConvertMessageToReplicaMessage(msg)

	// Serialize
	msgBytes, err := replicaMsg.Serialize()
	if err != nil {
		log.Printf("Node %d: Failed to serialize message: %v", n.config.ID, err)
		return
	}

	// Determine protocol type
	var mtype ProtocolType
	// CBR protocol uses CBR type
	if msg.Tag == "CBR" || msg.Tag == "CBR_PREPARE" || msg.Tag == "CBR_SHARE_REDUCE" || 
	   msg.Tag == "CBR_PROACTIVIZE" || msg.Tag == "CBR_SHARE_DIST" {
		mtype = CBR
	} else {
		// Default to CBR for unknown tags
		mtype = CBR
	}

	// Send to specific node using the client factory
	SendToNode(msgBytes, int64(to), mtype, n.addrs[to], CreateGRPCClient)
}

// WaitForFullConnection waits for full network connection (gRPC doesn't need this, but we keep the interface)
func (n *NetworkGRPC) WaitForFullConnection() {
	log.Printf("Node %d: gRPC network doesn't require explicit connection waiting", n.config.ID)
	// gRPC connections are established on-demand, so we just wait a bit
	time.Sleep(1 * time.Second)
}

// WaitForConnections waits for connections (kept for compatibility)
func (n *NetworkGRPC) WaitForConnections(timeout time.Duration) bool {
	log.Printf("Node %d: gRPC network doesn't require explicit connection waiting (timeout=%v)", n.config.ID, timeout)
	time.Sleep(1 * time.Second)
	return true
}

